from dna.Parameters import *
from dna.PolymorphismTables import PolymorphismTables
import dna.SiteFrequencySpectrum as sfs
import dna.infiniteSitesTools as ist
import numpy.matlib as np
import scipy.stats.mstats as mstats
import matplotlib.pyplot as plt

#test

NUM_OF_SIMULATIONS = 100
NUM_OF_CHROMOSOMES = 10

N = 1e4
MU = 1e-8
LENGTH = 25000
THETA = 4.0 * N * MU * LENGTH

p = Population(NUM_OF_CHROMOSOMES, N, LENGTH, MU, 0.0)
print p
print

d = Demography(N)
ep = d.addEpoch(100, 1e7, 0.0)
ep.setIterationParameters('time', 0.01, 1.2, 10.0, True)

td = ist.TajimasD()

print '\t  numM\t numV\t avgTD\t avgTDvar'

outData = []
t = []
for e in ep:
	pt = PolymorphismTables()
	msOut = pt.runMS( p.getMSString(NUM_OF_SIMULATIONS) + d.getMSString() )
	pt.readMS(None, msOut)

	mySFS = sfs.SiteFrequencySpectrum(pt)
	avgNum = td.calculateAllNumerator( mySFS )
	avgD = td.calculateAll(mySFS)
	avgDv = td.calculateAllVariance(mySFS)

	a = [ mstats.tmean( avgNum ) / mstats.var(avgNum), mstats.var( avgNum ), mstats.tmean(avgD), mstats.tmean(avgDv)] 
	outData.append(a)
	t.append(ep.myTime)
	print 't:%1.3f\t %1.3f\t %1.3f\t %1.3f\t %1.3f' % (ep.myTime, a[0], a[1], a[2], a[3] )
	pt = None

print
data = np.matrix(outData)
print data
print 
print t

plt.semilogx( t, np.ravel(data[:,0]), 'b', label='Num')
plt.semilogx( t, np.ravel(data[:,2]), 'r', label='TajD')
plt.semilogx( t, np.ravel(data[:,1]), 'b--', label='NumVar')
plt.semilogx( t, np.ravel(data[:,3]), 'r--', label='TajDVar') 
plt.legend(loc=1)
plt.ylabel('Value')
plt.xlabel('Time of event (in coalescent time)')
plt.title('CC model with 1000-fold reduction')
plt.show()
